import javax.swing.*;
import java.util.*;

/**
 * Created with IntelliJ IDEA.
 * Description:
 * User: 23735
 * Date: 2023-08-24
 * Time: 22:09
 */

class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode() {}
    TreeNode(int val) { this.val = val; }
    TreeNode(int val, TreeNode left, TreeNode right) {
        this.val = val;
        this.left = left;
        this.right = right;
    }
}

public class Solution {

    /**
     *  中序遍历非递归
     */
    public int kthSmallest1(TreeNode root, int k) {
        Stack<TreeNode> stack = new Stack<>();
        // 注意是怎么非递归的
        while (root != null || !stack.isEmpty()) {
            // 左子树
            while (root != null) {
                stack.push(root);
                root = root.left;
            }
            root = stack.pop();
            k--;
            if (k == 0) {
                return root.val;
            }
            root = root.right;
        }
        return -1;
    }

    /**
     *  记录以当前节点为根节点的节点数目
     */
    public int kthSmallest2(TreeNode root, int k) {
        BST bst = new BST(root);
        return bst.kthSmallest(k);
    }

}

class BST{
    TreeNode root;
    Map<TreeNode, Integer> nodeNum;

    public BST(TreeNode root) {
        this.root = root;
        this.nodeNum = new HashMap<>();
        getNodeNum(root);
    }

    public int kthSmallest(int k) {
        TreeNode node = root;
        while (node != null) {
            int leftNum = nodeNum.get(node.left);
            // 左边数目不够 k-1 个, 那么 kth 一定在右子树
            if (leftNum < k-1) {
                node = node.right;
                k = k - leftNum - 1;
            } else if (leftNum == k-1) {
                // 左边数目正好 k-1 个, 那么 当前节点就是 kth
                break;
            } else {
                // 左边数目大于 k-1 个, 那么 kth 一定在左子树
                node = node.left;
                // 注意这里 k 不能 --
                // k--;
            }
        }
        return node.val;
    }

    /**
     *  计算以当前节点为根节点的节点数目
     */
    private int getNodeNum (TreeNode node) {
        if (node == null) {
            nodeNum.put(node, 0);
            return 0;
        }
        nodeNum.put(node, (1+getNodeNum(node.left)+getNodeNum(node.right)));
        return nodeNum.get(node);
    }
}
